#include "pcap.hpp"
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <string>
#include "log.hpp"

bool PcapPaser::open()
{
    m_fd = ::open(m_file_path.c_str(), O_RDONLY);
    if (m_fd == -1)
    {
        return false;
    }

    // mmap all file
    {
        if (fstat(m_fd, &m_file_info) == -1)
        {
            return false;
        }
        m_file_content = ::mmap(nullptr, m_file_info.st_size,
                                PROT_READ, MAP_PRIVATE, m_fd, 0);
        if (m_file_content == MAP_FAILED)
        {
            return false;
        }
    }

    PcapFileHeader *file_header = (PcapFileHeader *)m_file_content;

    if (file_header->magic != 0xa1b2c3d4)
    {
        return false;
    }

    m_current = (uint8_t *)m_file_content + sizeof(PcapFileHeader);
    m_end = (uint8_t *)m_file_content + m_file_info.st_size;

    LOG("magic: %x\tversion: %d.%d\n", file_header->magic, file_header->version_major, file_header->version_minor);

    return true;
}

uint8_t *PcapPaser::getIP(const std::string src,
                          const std::string dest,
                          int protocol)
{
    PcapPacketHeader *pcap_header = (PcapPacketHeader *)m_current;
    EthernetFrameHeader *ethernet_header = (EthernetFrameHeader *)(m_current + sizeof(PcapPacketHeader));
    Ipv4PacketHeader *ipv4_header = (Ipv4PacketHeader *)(m_current + sizeof(PcapPacketHeader) + sizeof(EthernetFrameHeader));

    m_current += pcap_header->caplen + sizeof(PcapPacketHeader);

    if (ntohs(ethernet_header->type) != 0x0800) // 不是IP报文
    {
        LOG("not ip packet\n");
        return nullptr;
    }


    // IP字符串转换为uint32_t
    uint32_t src_ip = inet_addr(src.c_str());
    uint32_t dest_ip = inet_addr(dest.c_str());
    if (src_ip == INADDR_NONE || dest_ip == INADDR_NONE)
    {
        LOG("invalid ip address\n");
        return nullptr;
    }

    // 检查IP报文的源IP和目的IP是否匹配
    if (src_ip != ipv4_header->src_ip ||
        dest_ip != ipv4_header->dst_ip)
    {
        // LOG("ip address not match\n");
        return nullptr;
    }

    // 检查IP报文的协议是否匹配
    if (ipv4_header->protocol != protocol)
    {
        // LOG("protocol not match\n");
        return nullptr;
    }

    return (uint8_t *)ipv4_header;
}

uint8_t *PcapPaser::getTCP(const std::string src,
                           const std::string dest,
                           size_t* tcp_data_len)
{
    Ipv4PacketHeader *ip = (Ipv4PacketHeader*)getIP(src, dest, 6);
    if (ip == nullptr)
    {
        // LOG("not tcp packet\n");
        return nullptr;
    }
    size_t len = ntohs(((Ipv4PacketHeader*)ip)->total_length);
    size_t ip_head_len = (((Ipv4PacketHeader*)ip)->header_length) * 4;

    TcpPacketHeader* tcp_header = (TcpPacketHeader*)((uint8_t*)ip + ip_head_len);
    *tcp_data_len = len - ip_head_len - tcp_header->header_length * 4;

    return (uint8_t*)tcp_header;
}

uint8_t *PcapPaser::getUDP(const std::string src,
                           const std::string dest,
                           size_t *udp_data_len)
{
    uint8_t *ip = getIP(src, dest, 17);
    if (ip == nullptr)
    {
        LOG("not udp packet\n");
        return nullptr;
    }

    size_t len = ntohs(((Ipv4PacketHeader *)ip)->total_length);
    size_t ip_head_len = (((Ipv4PacketHeader *)ip)->header_length) * 4;
    UdpPacketHeader* udp_header = (UdpPacketHeader *)((uint8_t *)ip + ip_head_len);

    *udp_data_len = ntohs(udp_header->length) - sizeof(UdpPacketHeader);
    return (uint8_t*)udp_header;
}
